# Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional, Union

from mmengine import MessageHub
from mmengine.dist import get_rank
from mmengine.hooks import Hook

from xtuner.utils.device import get_device

DATA_BATCH = Optional[Union[dict, tuple, list]]


class VarlenAttnArgsToMessageHubHook(Hook):
    def before_train_iter(
        self, runner, batch_idx: int, data_batch: dict = None
    ) -> None:
        rank = get_rank()
        message_hub = MessageHub.get_instance("varlen_attn_args")

        assert "data" in data_batch.keys()
        data = data_batch["data"]

        cumulative_len = data.pop("cumulative_len")
        assert len(cumulative_len) == 1
        cumulative_len = cumulative_len[0].to(get_device())
        message_hub.update_info(f"cumulative_len_rank_{rank}", cumulative_len)

        max_seqlen = data.pop("max_seqlen")
        message_hub.update_info(f"max_seqlen_rank_{rank}", max_seqlen)

    def after_train_iter(
        self,
        runner,
        batch_idx: int,
        data_batch: DATA_BATCH = None,
        outputs: Optional[dict] = None,
    ) -> None:
        rank = get_rank()
        message_hub = MessageHub.get_instance("varlen_attn_args")
        message_hub.update_info(f"cumulative_len_rank_{rank}", None)
        message_hub.update_info(f"max_seqlen_rank_{rank}", None)

    def before_val_iter(
        self, runner, batch_idx: int, data_batch: DATA_BATCH = None
    ) -> None:
        """All subclasses should override this method, if they need any
        operations before each validation iteration.

        Args:
            runner (Runner): The runner of the validation process.
            batch_idx (int): The index of the current batch in the val loop.
            data_batch (dict, optional): Data from dataloader.
                Defaults to None.
        """
        rank = get_rank()
        message_hub = MessageHub.get_instance("varlen_attn_args")

        assert "data" in data_batch.keys()
        data = data_batch["data"]

        cumulative_len = data.pop("cumulative_len")
        assert len(cumulative_len) == 1
        cumulative_len = cumulative_len[0].to(get_device())
        message_hub.update_info(f"cumulative_len_rank_{rank}", cumulative_len)

        max_seqlen = data.pop("max_seqlen")
        message_hub.update_info(f"max_seqlen_rank_{rank}", max_seqlen)

    def after_val_iter(self, runner, batch_idx, data_batch=None, outputs=None) -> None:
        """All subclasses should override this method, if they need any
        operations after each validation iteration.

        Args:
            runner (Runner): The runner of the validation process.
            batch_idx (int): The index of the current batch in the val loop.
            data_batch (dict or tuple or list, optional): Data from dataloader.
            outputs (Sequence, optional): Outputs from model.
        """
        rank = get_rank()
        message_hub = MessageHub.get_instance("varlen_attn_args")
        message_hub.update_info(f"cumulative_len_rank_{rank}", None)
        message_hub.update_info(f"max_seqlen_rank_{rank}", None)
